This potentially violates the user's cpu_bytes_to_use, allocating more than the user allowed.
I think we want to add an alignment classvar in cpu/spec.py, which will be overrided in tiering/spec.py:

  # CPUOffloadingSpec
  class CPUOffloadingSpec(OffloadingSpec):
      CPU_PAGE_SIZE_ALIGNMENT = 1

      def __init__(self, vllm_config, kv_cache_config):
          ...
          kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
          self.cpu_page_size_per_worker = round_up(
              kv_bytes_per_offloaded_block // world_size,
              self.CPU_PAGE_SIZE_ALIGNMENT,
          )
          self.num_blocks = (
              int(cpu_bytes_to_use) // (self.cpu_page_size_per_worker * world_size)
              if self.cpu_page_size_per_worker > 0
              else 0
          )
          ...

  # TieringOffloadingSpec
  class TieringOffloadingSpec(CPUOffloadingSpec):
      CPU_PAGE_SIZE_ALIGNMENT = SharedOffloadRegion.PAGE_SIZE_ALIGNMENT
      ...

Nice catch. I agree that cpu_bytes_to_use should be respected.

Concern:
Expanding individual cpu_page_size_per_worker looks like it'll break a invariants and introduce some hard-to-catch bugs. i.e. We are going from,

  B0 |<--- B0 W0 ---><--- B0 W1 ---><--- B0 W2 --->| 
  B1 |<--- B1 W0 ---><--- B1 W1 ---><--- B1 W2 --->| 
  B2 |<--- B2 W0 ---><--- B2 W1 ---><--- B2 W2 --->| 
  ...
 
where Bi - Block i ; Wj - Worker j

to

  B0 |<--- B0 W0 ---***pad***><--- B0 W1 ---***pad***><--- B0 W2 ---***pad***>| 
  B1 |<--- B1 W0 ---***pad***><--- B1 W1 ---***pad***><--- B1 W2 ---***pad***>| 
  B2 |<--- B2 W0 ---***pad***><--- B2 W1 ---***pad****><--- B2 W2 ---***pad***>| 
  ...
 
where Bi - Block i ; Wj - Worker j

One example is the assert in cpu <-> gpu transfer
assert cpu_page_size == gpu_page_size * block_size_factor

Instead I propose doing,

  B0 |<--- B0 W0 ---><--- B0 W1 ---><--- B0 W2 --->***pad***| 
  B1 |<--- B1 W0 ---><--- B1 W1 ---><--- B1 W2 --->***pad***| 
  B2 |<--- B2 W0 ---><--- B2 W1 ---><--- B2 W2 --->***pad***| 
  ...

which is a looser constraint and can be handled directly in SharedOffloadRegion. Respecting cpu_bytes_to_use can be handled by,

• Allocating less num_blocks in CPUOffloadingSpec when padding is involved (communicated via CPU_PAGE_SIZE_ALIGNMENT classvar or a constructor arg)
• passing cpu_bytes_to_use directly to SharedOffloadRegion.
• And introducing padding in SharedOffloadRegion that respects alignment and cpu_bytes_to_use

Assuming we move the alignment code further up to spec.py, let's go back to a simple tensor allocation here, but using an aligned page size.

Keeping this as the new set of changes still updates SharedOffloadRegion directly. PTAL. Thanks 🙌

Why do we need this assert?

This assert is making sure that the num_blocks is not 0. It indicates that the user should increase cpu_bytes_to_use to run with CPU offloading enabled.

Let's use a classvar instead of introducing a new init param

Updated to introduce a classvar.
I am a bit uncomfortable with the semantic difference between CPUOffloadingSpec.BLOCK_SIZE_ALIGNMENT vs self.BLOCK_SIZE_ALIGNMENT seems like it increases the surface for bugs. what do you think ?

Same as the previous review: Why do we need this assert?

I replied in the previous review - #43689 (comment)

IMO This is a bit far fetched.
And if have a single block, what is it good for?

I have reverted the assert.
but, I believe it is better to be defensive in this case - for example I think having zero num_blocks will fail mmap in SharedOffloadRegion because num_blocks is zero. What do you think ? maybe we should handle it elsewhere.